#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# # author : cypro666
# # date   : 2015.08.01
# # wrapper of K-Means method in sklearn
import sys, json, time
import numpy as np
from threading import Thread
from sklearn import cluster

from magic3.utils import Timer
from magic3.filesystem import *
add_sys_path(grand_dir(__file__))

from skt.base import MethodBase
from skt.utils import _normalize, make_options


class KMeans(MethodBase):
    def __init__(self, parameters={}):
        super().__init__(parameters)

    @property
    def name(self): return 'KMeans'

    def output(self, model):
        self.Log('output')
        super().output()

        extra = {
            "centers" : [list(vec) for vec in model.cluster_centers_],
            "labels"  : [int(i) for i in model.labels_],
            "inertia" : float(model.inertia_)
        }

        fn = self.json_name()
        json.dump(extra, open(fn, 'w'), indent=4)

    def parallel(self):
        if self._param['n_clusters'] > 10:
            return 4
        else:
            return 1

    def execute(self):
        self.read_input()
        self._normalize(False, '1')

        if self._param['tol'] < 1.0E-5:
            self._param['tol'] = 1.0E-5

        if self._param['n_clusters'] < 2:
            self._param['n_clusters'] = 2

        if self._param['max_iter'] < 10:
            self._param['max_iter'] = 10

        self.save_parameters()

        model = cluster.KMeans(init='k-means++',
                               copy_x=False,
                               n_jobs=self.parallel(),
                               n_clusters=self._param['n_clusters'],
                               n_init=self._param['n_init'],
                               tol=self._param['tol'],
                               max_iter=self._param['max_iter'],
                               precompute_distances=True,
                               random_state=True,
                               verbose=False)

        if self._testing is not None:
            model.fit(self._train)
            self._results = model.predict(self._testing)
        else:
            self._results = model.fit_predict(self._train)

        self.output(model)

    def run(self, timeout):
        if self.parallel() == 1:
            t = Thread(target=self.execute)
            t.start()
            t.join(timeout)
            if t.is_alive():
                self.Log('timeout!')
            self.Log('exit')
        else:
            self.execute()


if __name__ == '__main__':
    opts = [('train_file', 'str', []),
            ('results_file', 'str', []),
            ('log_file', 'str', []),
            ('n_clusters', 'int', []),
            ('n_init', 'int', []),
            ('tol', 'float', []),
            ('max_iter', 'int', [])]

    parameters = make_options(opts)

    timer = Timer()
    KMeans(parameters).run(6000)

    if __debug__:
        print('elapsed:', timer)

